import threading
from MyTT import EMA
import numpy as np
import pandas as pd
from tqsdk import TqApi
from db_manager import connect_db, get_monitored_products  # 新增导入
import time

class OpenCloseStrategy:
    def __init__(self, api: TqApi, symbols, lots=1):  # 新增symbols参数
        self.api = api
        self.symbols = symbols  # 新增字段
        self.lots = lots
        self.running = False
        self.threads = []
        # 移除数据库连接（改由外部传入symbols）
    
    def start(self):
        """启动策略"""
        self.running = True
        # 修改为遍历传入的symbols
        for symbol in self.symbols:
            thread = threading.Thread(target=self._process_product, args=(symbol,))
            self.threads.append(thread)
            thread.start()
    
    def _process_product(self, symbol):  # 修改参数直接接收symbol
        """单个品种处理线程"""
        klines = self.api.get_kline_serial(symbol, 300, 500) # 改为从api获取k线数据,300秒即5分钟K线图,500个K线
        
        while self.running:
            print(symbol+" 正在运行。。。")
            # 新K线推送时
            if self.api.is_changing(klines.iloc[-1], "datetime"): 
                if not klines.close.iloc[-1]: # 如果最后一根K线的收盘价为空，跳过
                    print("K线已变化")
                    # 计算技术指标
                    q = (3*klines.close + klines.low + klines.open + klines.high) / 6
                    ma_values = self._calculate_ma(q)
                    short_ma = ma_values['ma']
                    ema_7 = ma_values['ema']
                    
                    # 获取当前持仓
                    position = self.api.get_position(symbol)
                    
                    # 修改后的交易逻辑
                    # 平全部空仓
                    if short_ma[-2] > ema_7[-2]:  
                        if position.pos_short > 0:
                            self.api.insert_order(symbol, direction="BUY", offset="CLOSE", 
                                            volume=position.pos_short)
                    # 开多仓
                    if position.pos_long == 0 and position.pos_short == 0 and short_ma[-3] < ema_7[-3] and short_ma[-2] > ema_7[-2]:
                        self.api.insert_order(symbol, direction="BUY", offset="OPEN", 
                                        volume=self.lots)
                        
                    # 平全部多仓
                    if short_ma[-2] < ema_7[-2]:  
                        if position.pos_long > 0:
                            self.api.insert_order(symbol, direction="SELL", offset="CLOSE", 
                                            volume=position.pos_long)
                    # 开空仓
                    if position.pos_long == 0 and position.pos_short == 0 and short_ma[-3] > ema_7[-3] and short_ma[-2] < ema_7[-2]:
                        self.api.insert_order(symbol, direction="SELL", offset="OPEN", 
                                        volume=self.lots)
                    # 检查是否需要更新K线数据
                    # if klines.iloc[-1].datetime != self.api.get_kline_serial(symbol, 300, 500).iloc[-1].datetime:
                    #     klines = self.api.get_kline_serial(symbol, 300, 500)  # 更新K线数据
            #deadline = time.time() + 0.5
            #self.api.wait_update(deadline=deadline)  # 等待0.5秒 
            time.sleep(1)

    
    def _calculate_ma(self, q):
        """自定义权重计算+MyTT优化"""
      
        # 保持原有的权重计算方式
        weights = np.arange(26, 0, -1)
        ma = []
        q_array = np.array(q)
        for i in range(len(q_array)):
            if i < 26:
                ma.append(np.nan)
                continue
            window = q_array[i-26:i+1]
            ma.append(np.dot(window, weights) / 351)
        
        # 使用MyTT的EMA计算
        ema_7 = EMA(pd.Series(ma), 7)
        
        return {'ma': ma, 'ema': ema_7.tolist()}
    
    def stop(self):
        """停止策略"""
        self.running = False
        for thread in self.threads:
            thread.join()

